{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimize Llama2 Model on Azure Machine Learning\n",
    "\n",
    "In this tutorial, we will optimize Llama2 model, and leverage AML (Azure Machine Learning) compute resources to run Olive Pass, while using your local device as the target to evaluate both the input model and output model.\n",
    "\n",
    "we will apply [QLoRA](https://github.com/artidoro/qlora), [ONNX conversion](https://huggingface.co/docs/transformers/serialization) and [ONNXRuntime transformers optimization](https://onnxruntime.ai/docs/performance/transformers-optimization.html) to optimize the original LLaMA 2 model, and evaluate the model performance by latency metric.\n",
    "\n",
    "## Prerequisites\n",
    "### installing Olive\n",
    "To get started with Olive, install it using the following command:\n",
    "```\n",
    "pip install git+https://github.com/microsoft/Olive#egg=olive-ai[azureml]\n",
    "```\n",
    "\n",
    "### Preparing AML compute\n",
    "Ensure your AML compute is set up and ready. If you haven't donw this yet, please refer to the [instructions](https://learn.microsoft.com/en-us/azure/machine-learning/how-to-create-compute-instance?view=azureml-api-2&tabs=azure-studio) for creating an AML compute instance.\n",
    "\n",
    "### Attaching your local machine to the AML workspace\n",
    "In this tutorial, we will use the local machine as target to run the model evaluation. Please follow [this guide](https://microsoft.github.io/Olive/tutorials/azure_arc.html) to attach your local machine to AML workspace.\n",
    "\n",
    "### Logging in to Azure with Azure CLI\n",
    "Authenticate and log in through your browser with the `az login` command.\n",
    "\n",
    "### Setting up Huggingface token in keyvault.\n",
    "To access Huggingface, please follow [this guide](https://microsoft.github.io/Olive/features/huggingface_model_optimization.html#azureml-system) to setup your Huggingface token."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Olive workflow\n",
    "### Json configuration file\n",
    "\n",
    "#### Azure ML client\n",
    "Since Olive will run Pass on your AML compute, add `azureml_client` section to the config with your workspace info:\n",
    "```\n",
    "\"azureml_client\": {\n",
    "    \"subscription_id\": \"<subscription_id>\",\n",
    "    \"resource_group\": \"<resource_group>\",\n",
    "    \"workspace_name\": \"<workspace_name>\",\n",
    "    \"keyvault_name\": \"<my_keyvault_name>\"\n",
    "}\n",
    "```\n",
    "\n",
    "#### Input model\n",
    "In this tutorial, we will use Azure Machine Learning Llama2 curated model. The input model will be automatically downloaded from the [Azure Model catalog](https://ml.azure.com/models/Llama-2-7b/version/13/catalog/registry/azureml-meta):\n",
    "```\n",
    "\"input_model\":{\n",
    "    \"type\": \"PyTorchModel\",\n",
    "    \"config\": {\n",
    "        \"model_path\": {\n",
    "            \"type\": \"azureml_registry_model\",\n",
    "            \"config\": {\n",
    "                \"name\": \"Llama-2-7b\",\n",
    "                \"registry_name\": \"azureml-meta\",\n",
    "                \"version\": \"13\"\n",
    "            }\n",
    "        },\n",
    "        \"model_file_format\": \"PyTorch.MLflow\",\n",
    "        \"hf_config\": {\n",
    "            \"model_name\": \"meta-llama/Llama-2-7b-hf\",\n",
    "            \"task\": \"text-generation\"\n",
    "        }\n",
    "    }\n",
    "}\n",
    "```\n",
    "\n",
    "#### System\n",
    "There are 2 systems defined here. The first is the compute resource designated for running the Pass. The second is the cluster on your loacal machine that previously had been attached to your AML workspace:\n",
    "```\n",
    "\"systems\": {\n",
    "    \"aml\": {\n",
    "        \"type\": \"AzureML\",\n",
    "        \"config\": {\n",
    "            \"accelerators\": [\"gpu\"],\n",
    "            \"hf_token\": true,\n",
    "            \"aml_compute\": \"<my_aml_compute>\",\n",
    "            \"aml_docker_config\": {\n",
    "                \"base_image\": \"mcr.microsoft.com/azureml/openmpi4.1.0-cuda11.6-cudnn8-ubuntu20.04\",\n",
    "                \"conda_file_path\": \"conda.yaml\"\n",
    "            }\n",
    "        }\n",
    "    },\n",
    "    \"azure_arc\": {\n",
    "        \"type\": \"AzureML\",\n",
    "        \"config\": {\n",
    "            \"accelerators\": [\"gpu\"],\n",
    "            \"hf_token\": true,\n",
    "            \"aml_compute\": \"<my_arc_compute>\",\n",
    "            \"aml_docker_config\": {\n",
    "                \"base_image\": \"mcr.microsoft.com/azureml/openmpi4.0-cuda11.6-cudnn8-ubuntu20.04\",\n",
    "                \"conda_file_path\": \"conda.yaml\"\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n",
    "```\n",
    "\n",
    "#### Data config\n",
    "Add a data section that will be used in QLoRA:\n",
    "```\n",
    "\"data_configs\": {\n",
    "    \"tiny_codes_train\": {\n",
    "        \"name\": \"tiny_codes_train\",\n",
    "        \"type\": \"HuggingfaceContainer\",\n",
    "        \"user_script\": \"user_script.py\",\n",
    "        \"components\": {\n",
    "            \"load_dataset\": {\n",
    "                \"type\": \"load_tiny_code_dataset\"\n",
    "            }\n",
    "        },\n",
    "        \"params_config\": {\n",
    "            \"data_name\": \"nampdn-ai/tiny-codes\",\n",
    "            \"split\": \"train\",                \n",
    "            \"component_kwargs\": {\n",
    "                \"load_dataset\": {\n",
    "                    \"language\": \"Python\"\n",
    "                },\n",
    "                \"pre_process_data\": {\n",
    "                    \"dataset_type\": \"corpus\",\n",
    "                    \"corpus_strategy\": \"join\",\n",
    "                    \"text_template\": \"### Question: {prompt} \\n### Answer: {response}\",\n",
    "                    \"source_max_len\": 1024\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n",
    "```\n",
    "\n",
    "You can find more details about how to configure data configs [here](https://microsoft.github.io/Olive/tutorials/configure_data.html). `user_script.py` is needed for this data config. Please check [user_script.py](./user_script.py) for details.\n",
    "\n",
    "\n",
    "#### Evaluators\n",
    "Add the latency metric to evaluator. We will use this metric to evaluate the output model:\n",
    "```\n",
    "\"evaluators\": {\n",
    "    \"common_evaluator\": {\n",
    "        \"metrics\": [\n",
    "            {\n",
    "                \"name\": \"latency\",\n",
    "                \"type\": \"latency\",\n",
    "                \"sub_types\": [{\"name\": \"avg\", \"goal\": {\"type\": \"percent-min-improvement\", \"value\": 10}}],\n",
    "                \"user_config\": {\n",
    "                    \"user_script\": \"user_script.py\",\n",
    "                    \"dataloader_func\": \"dataloader_func_for_merged\",\n",
    "                    \"func_kwargs\": {\n",
    "                        \"dataloader_func\": {\n",
    "                            \"model_id\": \"meta-llama/Llama-2-7b-hf\",\n",
    "                            \"past_seq_length\": 0,\n",
    "                            \"seq_length\": 8,\n",
    "                            \"max_seq_length\": 2048\n",
    "                        }\n",
    "                    },\n",
    "                    \"batch_size\": 2,\n",
    "                    \"io_bind\": true\n",
    "                }\n",
    "            }\n",
    "        ]\n",
    "    }\n",
    "}\n",
    "```\n",
    "\n",
    "`user_script.py` is needed for this evaluator. Please check [user_script.py](./user_script.py) for details.\n",
    "\n",
    "\n",
    "#### Passes\n",
    "Now we can add passes to the config file. We firstly add the QLoRA pass with training arguments:\n",
    "```\n",
    "\"passes\": {\n",
    "    \"qlora\": {\n",
    "        \"type\": \"QLoRA\",\n",
    "        \"config\": {\n",
    "            \"lora_dropout\": 0.1,\n",
    "            \"train_data_config\": \"tiny-codes-train\",\n",
    "            \"eval_dataset_size\": 1024,\n",
    "            \"training_args\": {\n",
    "                \"per_device_train_batch_size\": 2,\n",
    "                \"per_device_eval_batch_size\": 2,\n",
    "                \"gradient_accumulation_steps\": 1,\n",
    "                \"max_steps\": 100,\n",
    "                \"logging_steps\": 50,\n",
    "                \"save_steps\": 50,\n",
    "                \"evaluation_strategy\": \"steps\",\n",
    "                \"adam_beta2\": 0.999,\n",
    "                \"max_grad_norm\": 0.3,\n",
    "                \"load_best_model_at_end\": true\n",
    "            }\n",
    "        }\n",
    "    },\n",
    "}\n",
    "```\n",
    "More details about each parameters can be found at [QLoRA doc](https://github.com/artidoro/qlora)\n",
    "\n",
    "We will then convert the pytorch model to ONNX model by adding conversion pass:\n",
    "```\n",
    "\"convert\": {\n",
    "    \"type\": \"OptimumConversion\",\n",
    "    \"config\": {\n",
    "        \"target_opset\": 17,\n",
    "        \"save_as_external_data\": true,\n",
    "        \"all_tensors_to_one_file\": true,\n",
    "        \"torch_dtype\": \"float32\"\n",
    "    }\n",
    "}\n",
    "```\n",
    "\n",
    "Then we can apply ONNXRuntime transformers optimizations to this converted ONNX model:\n",
    "```\n",
    "\"optimize\": {\n",
    "    \"type\": \"OrtTransformersOptimization\",\n",
    "    \"config\": {\n",
    "        \"save_as_external_data\": true,\n",
    "        \"all_tensors_to_one_file\": true,\n",
    "        \"model_type\": \"gpt2\",\n",
    "        \"opt_level\": 0,\n",
    "        \"only_onnxruntime\": false,\n",
    "        \"keep_io_types\": false,\n",
    "        \"float16\": true,\n",
    "        \"use_gpu\": true,\n",
    "        \"optimization_options\": {\n",
    "            \"use_multi_head_attention\": false\n",
    "        }\n",
    "    }\n",
    "}\n",
    "```\n",
    "\n",
    "#### Engine\n",
    "Add the engine config as following:\n",
    "```\n",
    "\"engine\": {\n",
    "    \"log_severity_level\": 0,\n",
    "    \"evaluator\": \"common_evaluator\",\n",
    "    \"target\": \"azure_arc\",\n",
    "    \"host\": \"aml\",\n",
    "    \"execution_providers\": [\"CUDAExecutionProvider\"],\n",
    "    \"cache_dir\": \"cache\",\n",
    "    \"output_dir\" : \"models/llama2\"\n",
    "}\n",
    "```\n",
    "\n",
    "The configuration file can be found in [config.json](./config.json).\n",
    "\n",
    "### Run Olive Workflow\n",
    "Now you can run Olive with command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m olive.workflows.run --config config.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Output models and metrics\n",
    "\n",
    "Once the workflow run finishes, you will get a graph with Pass history and metrics results:\n",
    "\n",
    "\n",
    "| model_id | parent_model_id | from_pass | duration_sec | metrics |\n",
    "| ---- | ---- | ---- | ---- | ---- |\n",
    "| input_model | | | | {\"latency-avg\": 53.96777} |\n",
    "| 1_QLoRA | input_model | QLoRA | 4212.34 | |\n",
    "| 2_OnnxConversion | 1_QLoRA | OnnxConversion | 221.756 | |\n",
    "| 3_OrtTransformersOptimization | 2_OnnxConversion | OrtTransformersOptimization | 1231.44 | {\"latency-avg\": 27.68013} |\n",
    "\n",
    "The output model, which has an average latency of **27.68013**, shows a **48.7%** improvement in performance on your target machine in comparison to the original Llama2 model's average latency of **53.96777**. \n",
    "\n",
    "You can find the output ONNX model in the output folder `models/llama2`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nb",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
